import csv

from sklearn.feature_extraction import DictVectorizer
from sklearn import preprocessing
from sklearn import tree
from sklearn.externals.six import StringIO

myFile = open(r'AllElectronics.csv', 'rt')
reader = csv.reader(myFile)
headers = next(reader) # 特征

print(headers)

featureList = [] 
labelList = [] # class值

for row in reader:
    labelList.append(row[len(row)-1])
    rowDict = {}
    for i in range(1, len(row)-1):
        rowDict[headers[i]] = row[i]
    featureList.append(rowDict)

# DictVectorize对使用字典存储的数据进行特征抽取和向量化
vec = DictVectorizer()
# 转化成特征矩阵
dummyX = vec.fit_transform(featureList).toarray()

print('dummyX:\n' + str(dummyX))
# 输出各个维度的特征含义
print(vec.get_feature_names())
print('labelList:' + str(labelList))

# Vectorize class labels
lb = preprocessing.LabelBinarizer()
dummyY = lb.fit_transform(labelList)
# print('dummyY:' + str(dummyY))

# 创建决策树，参数选择代表属性选择度量的属性，该参数代表根据熵的大小，选取结点
clf = tree.DecisionTreeClassifier(criterion = 'entropy')
clf = clf.fit(dummyX, dummyY)
# print('clf:' + str(clf))

# 生成可视化树状图
# graphviz命令 dot -Tpdf inis.dot -o outpu.pdf
with open('allElectronicsInformationGainOri.dot', 'w') as f:
    # f = tree.export_graphviz(clf, out_file=f)
    f = tree.export_graphviz(clf, feature_names=vec.get_feature_names(), out_file=f)

oneRowX = dummyX[0, :]
print('oneRowX:' + str(oneRowX))

newRowX = oneRowX
newRowX[0] = 1
newRowX[2] = 0
print('newRowX:' + str([newRowX]))

predictedY = clf.predict(dummyX)
print(dummyY)
print("predictedY:" + str(predictedY))